"""
数据获取MCP服务
"""
from core.graph import build_graph
from fastmcp import FastMCP
from core.utils import agent_service
from models.response import Response
from loguru import logger
from typing import Dict

graph = build_graph()
mcp = FastMCP("DataFetch", host='0.0.0.0', port=8930)


@mcp.tool()
async def data_fetch(sessionId: str, query: str) -> Dict:
    """根据用户提问智能获取用户相关信息及相关数据表。
    Args:
        sessionId: string类型，会话ID
        query: string类型, 用户的提问
    Returns:
        msg: string类型，枚举值：success/error
        data: list类型，其中元素的类型为dict，包含三个必要字段：
            - type(数据类型)
            - description(数据描述)
            - data(具体数据)，当type = TableData时，data数据包含如下字段：
                  - query: 智能问数时的查询语句
                  - condition: 数据表的条件
                  - fields: 数据表字段的简略说明
                  - columns: 数据表字段的详细说明
                  - rows: 行数据
        error: 若msg=success, 则为None;  若msg=error, 则为具体的错误信息。
    """
    logger.info({
        "sessionId": sessionId,
        "question": query
    })
    try:

        # # 创建sessionId， 如果没有的话。
        # agent_service.call_agent(
        #     data={
        #         "agentId": '123456',
        #         "sessionId": sessionId,
        #         "question": query
        #     }
        # )

        # 获取数据并做预处理。
        result = await graph.ainvoke(
            input={
                'sessionId': sessionId,
                'query': query,
                'userinfo': None,
                'schemas': [],
                'rewrites': [],
                'datasets': [],
                'filtered': []
            },
            stream_mode=["custom"])
        resp = Response(data=result[-1][-1])
    except Exception as e:
        resp = Response(msg='error', error=str(e))
    return resp.model_dump()


if __name__ == "__main__":
    mcp.run(transport='sse')
